from modelscope import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, snapshot_download
import torch
import math
import os
from pathlib import Path
from datetime import datetime
import base64
from io import BytesIO
from typing import Optional, List, Union
import time
import uuid
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

# 设置HF_ENDPOINT使用国内镜像加速模型下载
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 减少显存碎片化
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# 启用推理模式，减少显存占用
torch.inference_mode()

# 定义请求和响应模型
class ImageGenerationRequest(BaseModel):
    prompt: str
    negative_prompt: Optional[str] = ""
    width: Optional[int] = 1024
    height: Optional[int] = 1024
    num_inference_steps: Optional[int] = 8
    seed: Optional[int] = 0

class ImageData(BaseModel):
    b64_json: str
    revised_prompt: str

class ImageGenerationResponse(BaseModel):
    created: int
    data: List[ImageData]

# 初始化模型
def initialize_model():
    print("=" * 50)
    print("开始初始化模型...")
    print("=" * 50)
    
    scheduler_config = {
        'base_image_seq_len': 256,
        'base_shift': math.log(3),
        'invert_sigmas': False,
        'max_image_seq_len': 8192,
        'max_shift': math.log(3),
        'num_train_timesteps': 1000,
        'shift': 1.0,
        'shift_terminal': None,
        'stochastic_sampling': False,
        'time_shift_type': 'exponential',
        'use_beta_sigmas': False,
        'use_dynamic_shifting': True,
        'use_exponential_sigmas': False,
        'use_karras_sigmas': False,
    }

    print("正在配置调度器...")
    scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
    print("调度器配置完成")

    # 检测可用GPU数量
    num_gpus = torch.cuda.device_count()
    print(f"检测到 {num_gpus} 个GPU设备")

    # 根据项目规范，对于DiffusionPipeline模型，我们使用max_memory参数进行显存分配
    # 而不是手动指定每层的设备映射
    if num_gpus > 1:
        print("检测到多个GPU，正在进行显存分配...")
        # 获取每个GPU的显存信息并计算可分配的显存量
        max_memory = {}
        for i in range(num_gpus):
            free_mem, total_mem = torch.cuda.mem_get_info(i)
            # 按70%的空闲显存计算分配量，同时确保不超过22GB
            allocated_mem = min(int(free_mem * 0.7), 22 * 1024**3, free_mem)
            max_memory[i] = allocated_mem
            print(f"GPU {i}: 分配 {(allocated_mem / 1024**3):.2f} GB 显存")
        
        # 加载模型并指定显存分配
        print("正在加载模型到多个GPU...")
        pipe = DiffusionPipeline.from_pretrained(
            'Qwen/Qwen-Image',
            scheduler=scheduler,
            torch_dtype=torch.bfloat16,
            max_memory=max_memory,  # 为每个GPU分配显存
        )
        print("多GPU模型加载完成")
    else:
        # 单GPU情况
        print("检测到单个GPU，正在加载模型...")
        pipe = DiffusionPipeline.from_pretrained(
            'Qwen/Qwen-Image',
            scheduler=scheduler,
            torch_dtype=torch.bfloat16,
        )
        pipe = pipe.to("cuda")
        print("单GPU模型加载完成")

    print(f"模型已分配到{num_gpus}个GPU设备上")
    print("=" * 50)
    print("模型初始化完成")
    print("=" * 50)
    return pipe

# 提前下载LoRA权重
def download_lora_weights():
    print("=" * 50)
    print("开始下载LoRA权重...")
    print("=" * 50)
    # 使用ModelScope的snapshot_download下载LoRA权重
    model_dir = snapshot_download('lightx2v/Qwen-Image-Lightning')
    print(f"LoRA权重已下载到: {model_dir}")
    
    # 查找.pt或.safetensors文件
    lora_files = list(Path(model_dir).glob("*.safetensors")) + list(Path(model_dir).glob("*.pt"))
    if not lora_files:
        raise FileNotFoundError("在下载的LoRA权重目录中未找到.safetensors或.pt文件")
    
    lora_file_path = lora_files[0]  # 使用找到的第一个文件
    print(f"使用LoRA文件: {lora_file_path}")
    print("=" * 50)
    print("LoRA权重下载完成")
    print("=" * 50)
    return str(lora_file_path)

# 图像生成函数
def generate_image(request: ImageGenerationRequest, pipe):
    """生成图像"""
    print("=" * 50)
    print("开始生成图像...")
    print("=" * 50)
    
    print(f"输入参数:")
    print(f"  - 提示词: {request.prompt}")
    print(f"  - 负面提示词: {request.negative_prompt}")
    print(f"  - 图像尺寸: {request.width}x{request.height}")
    print(f"  - 推理步数: {request.num_inference_steps}")
    print(f"  - 随机种子: {request.seed}")
    
    start_time = time.time()
    print("开始图像生成过程...")
    
    print("正在调用模型生成图像...")
    image = pipe(
        prompt=request.prompt,
        negative_prompt=request.negative_prompt,
        width=request.width,
        height=request.height,
        num_inference_steps=request.num_inference_steps,
        true_cfg_scale=1.0,
        generator=torch.manual_seed(request.seed)
    ).images[0]
    
    generation_time = time.time() - start_time
    print(f"图像生成完成，耗时: {generation_time:.2f} 秒")
    
    print("正在将图像转换为base64编码...")
    # 将图像转换为base64编码
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    print("图像转换完成")
    
    print("=" * 50)
    print("图像生成流程结束")
    print("=" * 50)
    
    return img_str

# 初始化FastAPI应用
app = FastAPI(
    title="Qwen-Image API",
    description="基于Qwen-Image模型的OpenAI兼容图像生成API",
    version="1.0.0"
)

# 全局模型实例
pipe = None

@app.on_event("startup")
async def startup_event():
    global pipe
    # 初始化模型
    pipe = initialize_model()
    
    # 下载LoRA权重
    lora_file_path = download_lora_weights()
    print("正在加载LoRA权重...")
    pipe.load_lora_weights(lora_file_path)
    print("LoRA权重加载完成")

@app.post("/v1/images/generations", response_model=ImageGenerationResponse)
async def create_image(request: ImageGenerationRequest):
    try:
        # 生成图像
        image_data = generate_image(request, pipe)
        
        # 构造响应
        response = ImageGenerationResponse(
            created=int(time.time()),
            data=[
                ImageData(
                    b64_json=image_data,
                    revised_prompt=request.prompt
                )
            ]
        )
        
        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8800)